package find;

import java.util.Arrays;

/**
 * 在一个无序的数组中找到第k小个数, 和前k小个数
 * 最优解（Best friend party）:MIT 5个大佬想出来的，主要优化了荷兰国旗中选择划分数这一步，避免了较差的划分情况
 * 时间复杂度:O(N)
 *
 * @author Liaorun
 */
@SuppressWarnings({"AlibabaLowerCamelCaseVariableNaming", "SpellCheckingInspection", "AlibabaClassNamingShouldBeCamel"})
public class BFPRT {

    public static void main(String[] args) {
        System.out.println(getMinKthByBFPRT(new int[]{5, 4, 3, 2, 2, 1, 0}, 7));
        System.out.println(Arrays.toString(getMinKNumsByBFPRT(new int[]{5, 4, 3, 2, 2, 1, 0}, 7)));
    }

    /**
     * 获取一个无序数组的前K小个数
     *
     * @param arr 数组
     * @param k   排序后的下标位置
     * @return 排序后的下标位置 <= k 的所有数值
     */
    public static int[] getMinKNumsByBFPRT(int[] arr, int k) {
        if (k < 1 || k > arr.length) {
            return arr;
        }

        int minKth = getMinKthByBFPRT(arr, k);

        int[] res = new int[k];
        int index = 0;


        for (int value : arr) {
            // 原数组中所有小于第k位置上的数全部copy到结果数组
            if (value < minKth) {
                res[index++] = value;
            }
        }


        for (; index < res.length; index++) {
            // 结果数组末尾空着的地方是，荷兰国旗中间的部分，全部填入第K个位置的值（划分值）
            res[index] = minKth;
        }

        return res;
    }

    /**
     * 获取一个无序数组第K小的数
     *
     * @param arr 数组
     * @param k   顺序
     * @return 排序后第K位置的数值
     */
    public static int getMinKthByBFPRT(int[] arr, int k) {
        // patition会修改原数组数值的位置，所有需要copy一份出来处理
        int[] copyArr = copyArray(arr);

        return select(copyArr, 0, copyArr.length - 1, k - 1);
    }

    private static int[] copyArray(int[] arr) {
        int[] res = new int[arr.length];

        System.arraycopy(arr, 0, res, 0, res.length);

        return res;
    }

    /**
     * 在arr[begin..end]的范围上，如果排序的话，i位置的数是谁
     *
     * @param arr   数组
     * @param begin 开始位置
     * @param end   结束位置
     * @param i     一定在begin..end范围上
     * @return 排序后第K个位置的值
     */
    private static int select(int[] arr, int begin, int end, int i) {
        if (begin == end) {
            return arr[begin];
        }

        // 分组 + 组内排序 + 每组的中位数组成新数组 + 选出新数组的上中位数
        int pivot = medianOfMedians(arr, begin, end);

        int[] pivotRange = partition(arr, begin, end, pivot);
        if (i >= pivotRange[0] && i <= pivotRange[1]) {
            // 需要找的数在国旗的中间部分，找到第k个数，直接返回
            return arr[i];
        } else if (i < pivotRange[0]) {
            // 需要找的数在国旗的中间部分左边，在左边去接着递归找
            return select(arr, begin, pivotRange[0] - 1, i);
        } else {
            // 需要找的数在国旗的中间部分右边，在右边去接着递归找
            return select(arr, pivotRange[1] + 1, end, i);
        }
    }

    private static int[] partition(int[] arr, int begin, int end, int pivot) {
        int small = begin - 1;
        int cur = begin;
        int big = end + 1;

        // 跳出条件：和大于区域相撞
        while (cur != big) {
            if (arr[cur] < pivot) {
                // 当前值小于划分值，和小于区域的下一个交换，cur++
                swap(arr, ++small, cur++);
            } else if (arr[cur] > pivot) {
                // 当前值小于划分值，和大于区域的前一个交换，cur不变
                swap(arr, cur, --big);
            } else {
                // 当前值等于划分值，看下一个
                cur++;
            }
        }
        return new int[]{small + 1, big - 1};
    }

    private static void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    /**
     * 给荷兰国旗问题选择一个合适的中间划分数
     * 分组 + 组内排序 + 每组的中位数组成新数组 + 选出新数组的上中位数
     *
     * @param arr   数组
     * @param begin 开始位置
     * @param end   结束位置
     * @return patition合适的划分值
     */
    private static int medianOfMedians(int[] arr, int begin, int end) {
        int num = end - begin + 1;
        int offset = num % 5 == 0 ? 0 : 1;
        int[] mArr = new int[num / 5 + offset];

        for (int i = 0; i < mArr.length; i++) {
            int beginI = begin + i * 5;
            int endI = beginI + 4;
            mArr[i] = getMedian(arr, beginI, Math.min(endI, end));
        }
        return select(mArr, 0, mArr.length - 1, mArr.length / 2);
    }

    /**
     * 获取数组的中位数
     *
     * @param arr   数组
     * @param begin 开始位置
     * @param end   结束位置
     * @return 中位数 or 上中位数
     */
    private static int getMedian(int[] arr, int begin, int end) {

        Arrays.sort(arr, begin, end);

        int sum = end + begin;
        int mid = (sum / 2) + (sum % 2);

        return arr[mid];
    }
}
